# $Id: tcp.py 9227e81e919a 2012/08/27 15:49:57 pgurumur $
# Copyright (c) 2009 Prabhu Gurumurthy  <pgurumur@gmail.com>
# 
# Permission to use, copy, modify, and distribute this software for any
# purpose with or without fee is hereby granted, provided that the above
# copyright notice and this permission notice appear in all copies.
#
# THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES
# WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
# MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR
# ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
# WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN
# ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF
# OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
#

__version__  = "$Revision: 9227e81e919a $"
__author__   = "$Author: pgurumur $"
__modified__ = "$Date: 2012-08-27 15:49:57Z $"

# essential imports
import select
import socket
import threading
import SocketServer

# Error class
class SocketError(Exception):
   pass

# tcp client class
class tcp(object):
   def __init__(self, Host = None, Port = 0):
      # init the error class
      self._err = SocketError()

      # Local defines
      self._sock = None
      self._host = None
      self._port = 0

      # socket list for getaddrinfo
      self._sockList = []

      if Host:
         self._host = Host
      else:
         self._error("no host defined!")

      if Port:
         try:
            self._port = int(Port)
         except ValueError, message:
            self._error(message)

      # Try using getaddrinfo, we raise SocketError, on errors
      try:
         self._sockList = socket.getaddrinfo(self._host, self._port,
               socket.AF_UNSPEC, socket.SOCK_STREAM)
      except socket.gaierror, (error, message):
         self._error("getaddrinfo: error %d, message: %s" %(error, message))

   # Function used when print is called on the class
   def __str__(self):
      retval = None
      if self._sock:
         retval = "tcp socket %d to host %s on port %d" %(
               self._sock.fileno(), self._host, self._port)

      return retval

   # Get the file descriptor, when I call int() on the socket
   def __int__(self):
      retval = -1
      if self._sock:
         retval = self._sock.fileno()

      return retval

   # error handling the function, we close the socket first and then 
   # raise the exception
   def _error(self, Message):
      self.Close()
      if Message:
         raise SocketError, Message

   # Close the socket, if it is open
   def Close(self):
      if self._sock:
         self._sock.close()

   # Does an actual connect on tcp socket
   def Connect(self):
      # mesg variable to be used when there is an error
      mesg = None
      for res in self._sockList:
         af, socktype, proto, canonname, sa = res
         try:
            self._sock = socket.socket(af, socktype, proto)
         except socket.error, (error, message):
            self._sock = None
            continue

         try:
            # non-blocking socket
            self._sock.setblocking(0)

            # connect
            self._sock.connect(sa)
         except socket.error, (error, message):
            # if we get EINPROGRESS, then pass, otherwise error out
            if error == 36:
               pass
            else:
               mesg = message
               self._sock.close()
               self._sock = None
               continue

         break

      if self._sock is None:
         # If socket is none, print why we cannot connect
         self._error("cannot connect to %s on tcp port %d: %s" %(
            self._host, self._port, mesg))
      else:
         # set socket to be blocking after we connect
         self._sock.setblocking(1)

   def Write(self, Data, Timeout = 0):
      # how many bytes have I sent out!
      sent = None
      if Data:
         try:
            # If timeout is present, use it to select on the socket to see
            # whether it is ready to send data out!
            # otherwise, wait forever, to select on the socket
            if Timeout:
               read, write, err = select.select([], [self._sock], [], Timeout)
            else:
               read, write, err = select.select([], [self._sock], [])
         except select.error, (error, message):
            self._error("select error: %d mesg: %s" %(error, message))

         # socket is ready to write, then send the data, if there is a socket
         # error, exit
         if self._sock in write:
            try:
               sent = self._sock.send(Data)
            except socket.error, (error, message):
               self._error("socket error: %s" %message)
      else:
         self._error("no data to send")

      # Assert to see whether we sent the data correctly
      assert(sent == len(Data))

      # return how much I sent
      return sent

   def Read(self, bufSize = 1024, Timeout = 0):
      if not bufSize:
         bufSize = 1024

      data = ""
      flag = True
      while flag:
         try:
            if Timeout:
               read, write, err = select.select([self._sock], [], [], Timeout)
            else:
               read, write, err = select.select([self._sock], [], [])
         except select.error, (error, message):
            self._error("select error: %d, mesg: %s" %(error, message))
            pass

         if self._sock in read:
            try:
               data += self._sock.recv(bufSize)
            except socket.error, (error, message):
               pass
         else:
            flag = False

      return data

class tcpClientHandler(SocketServer.StreamRequestHandler):
   def handle(self):
      self._threadId = threading.currentThread()

   def Read(self, bufSize = 1024, Timeout = 0):
      if not bufSize:
         bufSize = 1024

      data = ""
      flag = True
      while flag:
         try:
            if Timeout:
               read, write, err = select.select([self.request], [], [],
                     Timeout)
            else:
               read, write, err = select.select([self.request], [], [])
         except select.error, (error, message):
            raise SocketError, message

         if self.request in read:
            try:
               data = self.request.recv(bufSize)
            except socket.error, (error, message):
               pass
         else:
            flag = False

      return data

   def Write(self, Data, Timeout = 0):
      # how many bytes have I sent out!
      sent = None
      if Data:
         try:
            # If timeout is present, use it to select on the socket to see
            # whether it is ready to send data out!
            # otherwise, wait forever, to select on the socket
            if Timeout:
               read, write, err = select.select([], [self._sock], [], Timeout)
            else:
               read, write, err = select.select([], [self._sock], [])
         except select.error, (error, message):
            self._error("select error: %d mesg: %s" %(error, message))

         # socket is ready to write, then send the data, if there is a socket
         # error, exit
         if self._sock in write:
            try:
               sent = self._sock.send(Data)
            except socket.error, (error, message):
               self._error("socket error: %s" %message)
      else:
         self._error("no data to send")

      # Assert to see whether we sent the data correctly
      assert(sent == len(Data))

      # return how much I sent
      return sent

class ThreadedTCPServer(SocketServer.ThreadingMixIn, SocketServer.TCPServer):
   pass

class tcpServer:
   def __init__   (self, Port, Host = None):
      self._err = SocketError()
      self._port = 0
      self._host = None
      self._server = None
      self._thread = None

      if Port:
         try:
            self._port = int(Port)
         except ValueError, message:
            self._error(message)

      if Host:
         self._host = Host
      else:
         self._host = ""

   def Close(self):
      if self._server:
         self._server.socket.close()

   def _error(self, Message):
      self.Close()
      if Message:
         raise SocketError, message

   def Listen(self):
      self._server = ThreadTCPServer((self._host, self._port),
            tcpClientHandler)
      self._thread = threading.Thread(
            target = self._server.serve_forever)
      self._thread.setDaemon(True)
      self._thread.start()
